package com.shujia.spark.mllib

import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.regression.{LinearRegression, LinearRegressionModel}
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo2Line {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local")
      .appName("line")
      .getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._

    //读取原始数据
    val linesDF: DataFrame = spark.read
      .format("csv")
      .option("sep", ",")
      .schema("label DOUBLE, x STRING")
      .load("data/lines.txt")

    /**
     * 1、特征工程
     */
    //将特征转换成特征向量
    val toVector: UserDefinedFunction = udf((x: String) => {
      //将x转换成向量返回
      //Vectors.dense(x.toDouble)

      //切分成多个x，转换成double
      val array: Array[Double] = x
        .split("\\|")
        .map(i => i.toDouble)

      //将x转换成向量返回
      Vectors.dense(array)
    })

    //将数据转换成特征向量
    val dataDF: DataFrame = linesDF
      .select($"label", toVector($"x") as "features")


    /**
     * 2、选择算法，
     * 如果目标值是连续的选择回归
     * 如果目标值是离散的选择分类
     *
     * LinearRegression： 线性回归
     */

    val lr: LinearRegression = new LinearRegression()


    /**
     * 3、将数据带入算法训练模型
     */

    val model: LinearRegressionModel = lr.fit(dataDF)

    //查看训练结果
    println(model.intercept)
    println(model.coefficients)

    /**
     * 预测新的x
     */

    val y: Double = model.predict(Vectors.dense(100, 100, 100))
    println(y)

  }
}
